package com.cmnit.gatherdata.modules.utils

import java.io.File
import java.security.PrivilegedAction
import java.sql.SQLException

import com.cmnit.gatherdata.utils.ConfigurationManager
import org.apache.hadoop.hbase.client._
import org.apache.hadoop.hbase.{Cell, HBaseConfiguration, TableName}
import org.apache.hadoop.security.UserGroupInformation

/**
 * hbase工具类
 * 1.获取数据库链接
 * 2.根据rowkey查询
 * 3.根据rowkey插入
 *
 * @author chan
 * @version 1.2
 * @since 2022.3.3
 */
object HBaseUtil {

  /**
   * 获取Kerberos环境下的HBase连接
   *
   * @return
   */
  def getHBaseConn: Connection = {
    val configuration = HBaseConfiguration.create
    if (ConfigurationManager.getProperty("product.or.test") == "test") {
      val coreFile = HBaseUtil.getClass.getClassLoader.getResourceAsStream("core-site-test.xml")
      configuration.addResource(coreFile)
      val hdfsFile = HBaseUtil.getClass.getClassLoader.getResourceAsStream("hdfs-site-test.xml")
      configuration.addResource(hdfsFile)
      val hbaseFile = HBaseUtil.getClass.getClassLoader.getResourceAsStream("hbase-site-test.xml")
      configuration.addResource(hbaseFile)
    } else if (ConfigurationManager.getProperty("product.or.test") == "product") {
      val coreFile = HBaseUtil.getClass.getClassLoader.getResourceAsStream("core-site.xml")
      configuration.addResource(coreFile)
      val hdfsFile = HBaseUtil.getClass.getClassLoader.getResourceAsStream("hdfs-site.xml")
      configuration.addResource(hdfsFile)
      val hbaseFile = HBaseUtil.getClass.getClassLoader.getResourceAsStream("hbase-site.xml")
      configuration.addResource(hbaseFile)
    }

    UserGroupInformation.loginUserFromKeytab(ConfigurationManager.getProperty("kerberos.principal"), ConfigurationManager.getProperty("keytab.file"))
    UserGroupInformation.setConfiguration(configuration)

    val loginUser = UserGroupInformation.getLoginUser
    loginUser.doAs(new PrivilegedAction[Connection] {
      override def run(): Connection = ConnectionFactory.createConnection(configuration)
    })
    ConnectionFactory.createConnection(configuration)
  }

  /**
   * 获取非Kerberos环境的Connection
   *
   * @param confPath 参数路径
   * @return
   */
  def getNoKBHBaseCon(confPath: String): Connection = {

    val configuration = HBaseConfiguration.create
    val coreFile = new File(confPath + File.separator + "core-site.xml")
    if (!coreFile.exists()) {
      val in = HBaseUtil.getClass.getClassLoader.getResourceAsStream("hbase-conf/core-site.xml")
      configuration.addResource(in)
    }
    val hdfsFile = new File(confPath + File.separator + "hdfs-site.xml")
    if (!hdfsFile.exists()) {
      val in = HBaseUtil.getClass.getClassLoader.getResourceAsStream("hbase-conf/hdfs-site.xml")
      configuration.addResource(in)
    }
    val hbaseFile = new File(confPath + File.separator + "hbase-site.xml")
    if (!hbaseFile.exists()) {
      val in = HBaseUtil.getClass.getClassLoader.getResourceAsStream("hbase-conf/hbase-site.xml")
      configuration.addResource(in)
    }

    val connection = ConnectionFactory.createConnection(configuration)
    connection
  }

  /**
   * 根据rowkey进行查询，判断是否为空
   *
   * @param rowkey Rowkey
   * @return
   */
  def isEmptyByRowkey(tableName: String, rowkey: String, connect: Connection): Boolean = {
    var table: Table = null
    var result: Boolean = false
    try {
      table = connect.getTable(TableName.valueOf(tableName))
      // 获取结果集
      val cells = HBaseUtil.queryByRowkey(rowkey, table, connect)
      result = cells.isEmpty
    } catch {
      case e: Exception => println(e)
    }
    finally {
      HBaseUtil.close(null, table, null)
    }
    result
  }

  /**
   * 根据rowkey进行查询
   *
   * @param rowkey Rowkey
   * @param table  表连接
   * @param conn   jdbc连接
   * @return
   */
  def queryByRowkey(rowkey: String, table: Table, conn: Connection): Array[Cell] = {
    val result = table.get(new Get(rowkey.getBytes))
    result.rawCells
  }

  /**
   * 根据rowkey插入数据
   *
   * @param rowkey Rowkey
   * @param topic  topic
   */
  def upsertByRowkey(tableName: String, rowkey: String, topic: String, connect: Connection): Unit = {
    var table: Table = null
    try {
      table = connect.getTable(TableName.valueOf(tableName))
      val put = new Put(rowkey.getBytes)
      put.addColumn("info".getBytes, "id".getBytes, rowkey.getBytes)
      put.addColumn("info".getBytes, "topic".getBytes, topic.getBytes)
      table.put(put)
    } catch {
      case e: SQLException => e.printStackTrace()
    } finally {
      HBaseUtil.close(null, table, null)
    }
  }

  /**
   * 根据rowkey获取PUT对象
   *
   * @param rowkey rowkey
   * @param topic  topic
   * @return
   */
  def getPutByRowkey(rowkey: String, topic: String): Put = {
    val put = new Put(rowkey.getBytes)
    put.addColumn("info".getBytes, "id".getBytes, rowkey.getBytes)
    put.addColumn("info".getBytes, "topic".getBytes, topic.getBytes)
    put
  }

  /**
   * 关闭资源
   *
   * @param conn          jdbc连接
   * @param table         表连接
   * @param resultScanner 结果集
   */
  def close(conn: Connection, table: BufferedMutator, resultScanner: ResultScanner): Unit = {
    if (conn != null) try conn.close()
    catch {
      case e: SQLException => e.printStackTrace()
    }
    if (table != null) try table.close()
    catch {
      case e: SQLException => e.printStackTrace()
    }
    if (resultScanner != null) try conn.close()
    catch {
      case e: SQLException => e.printStackTrace()
    }
  }

  /**
   * 关闭资源
   *
   * @param conn          jdbc连接
   * @param table         表连接
   * @param resultScanner 结果集
   */
  def close(conn: Connection, table: Table, resultScanner: ResultScanner): Unit = {
    if (conn != null) try conn.close()
    catch {
      case e: SQLException => e.printStackTrace()
    }
    if (table != null) try table.close()
    catch {
      case e: SQLException => e.printStackTrace()
    }
    if (resultScanner != null) try conn.close()
    catch {
      case e: SQLException => e.printStackTrace()
    }
  }
}
